{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Complex Resonator Model\n",
    "\n",
    "This notebook shows how to fit the parameters of a complex resonator, \n",
    "using `lmfit.Model` and defining a custom `Model` class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pylab inline\n",
    "#%config InlineBackend.figure_format='retina'  # for hi-dpi displays"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import lmfit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(123) # make this notebook reproducible. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Following Khalil et al. (http://arxiv.org/abs/1108.3117), we can model the forward transmission of a microwave resonator with total quality factor $Q$, coupling quality factor $Q_e$, and resonant frequency $f_0$ using:\n",
    "\n",
    "$$ S_{21}(f) = 1-\\frac{Q\\,Q_e^{-1}}{1+2jQ(f-f_0)/f_0} $$\n",
    "\n",
    "$S_{21}$ is thus a complex function of a real frequency.\n",
    "\n",
    "By allowing $Q_e$ to be complex, this model can take into account mismatches in the input and output transmission impedances.\n",
    "\n",
    "Since `scipy.optimize` and `lmfit` require real parameters, we represent $Q_e$ as `Q_e_real + 1j*Q_e_imag`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def linear_resonator(f, f_0, Q, Q_e_real, Q_e_imag):\n",
    "    Q_e = Q_e_real + 1j*Q_e_imag\n",
    "    return (1 - (Q * Q_e**-1 / (1 + 2j * Q * (f - f_0) / f_0)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The standard practice of defining an `lmfit` model is as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ResonatorModel(lmfit.model.Model):\n",
    "    __doc__ = \"resonator model\" + lmfit.models.COMMON_DOC\n",
    "    \n",
    "    def __init__(self, *args, **kwargs):\n",
    "        # pass in the defining equation so the user doesn't have to later.\n",
    "        super(ResonatorModel, self).__init__(linear_resonator, *args, **kwargs)\n",
    "        \n",
    "        self.set_param_hint('Q', min=0)  # Enforce Q is positive\n",
    "    \n",
    "    def guess(self, data, f=None, **kwargs):\n",
    "        verbose = kwargs.pop('verbose', None)\n",
    "        if f is None:\n",
    "            return\n",
    "        argmin_s21 = np.abs(data).argmin()\n",
    "        fmin = f.min()\n",
    "        fmax = f.max()\n",
    "        f_0_guess = f[argmin_s21]  # guess that the resonance is the lowest point\n",
    "        Q_min = 0.1 * (f_0_guess/(fmax-fmin))  # assume the user isn't trying to fit just a small part of a resonance curve.\n",
    "        delta_f = np.diff(f)  # assume f is sorted\n",
    "        min_delta_f = delta_f[delta_f > 0].min()\n",
    "        Q_max = f_0_guess/min_delta_f  # assume data actually samples the resonance reasonably\n",
    "        Q_guess = np.sqrt(Q_min*Q_max)  # geometric mean, why not?\n",
    "        Q_e_real_guess = Q_guess/(1-np.abs(data[argmin_s21]))\n",
    "        if verbose:\n",
    "            print(\"fmin=\", fmin, \"fmax=\", fmax, \"f_0_guess=\", f_0_guess)\n",
    "            print(\"Qmin=\", Q_min, \"Q_max=\", Q_max, \"Q_guess=\", Q_guess, \"Q_e_real_guess=\", Q_e_real_guess)\n",
    "        params = self.make_params(Q=Q_guess, Q_e_real=Q_e_real_guess, Q_e_imag=0, f_0=f_0_guess)\n",
    "        params['%sQ' % self.prefix].set(min=Q_min, max=Q_max)\n",
    "        params['%sf_0' % self.prefix].set(min=fmin, max=fmax)\n",
    "        return lmfit.models.update_param_vals(params, self.prefix, **kwargs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's use the model to generate some fake data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "resonator = ResonatorModel()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "true_params = resonator.make_params(f_0=100, Q=10000, Q_e_real=9000, Q_e_imag=-9000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f = np.linspace(99.95, 100.05, 100)\n",
    "true_s21 = resonator.eval(params=true_params,f=f)\n",
    "noise_scale = 0.02\n",
    "measured_s21 = true_s21 + noise_scale*(np.random.randn(100) + 1j*np.random.randn(100))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot(f, 20*np.log10(np.abs(measured_s21)))\n",
    "ylabel('|S21| (dB)')\n",
    "xlabel('MHz')\n",
    "title('simulated measurement');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Try out the guess method we added:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "guess = resonator.guess(measured_s21, f=f, verbose=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And now fit the data using the guess as a starting point:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result = resonator.fit(measured_s21, params=guess, f=f, verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(result.fit_report())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result.params.pretty_print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we'll make some plots of the data and fit. Define a convenience function for plotting complex quantities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_ri(data, *args, **kwargs):\n",
    "    plot(data.real, data.imag, *args, **kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fit_s21 = resonator.eval(params=result.params, f=f)\n",
    "guess_s21 = resonator.eval(params=guess, f=f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_ri(measured_s21, '.')\n",
    "plot_ri(fit_s21, 'r.-')\n",
    "plot_ri(guess_s21, 'k--')\n",
    "xlabel('Re(S21)')\n",
    "ylabel('Im(S21)')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot(f, 20*np.log10(np.abs(measured_s21)), '.')\n",
    "plot(f, 20*np.log10(np.abs(fit_s21)), 'r.-')\n",
    "plot(f, 20*np.log10(np.abs(guess_s21)), 'k--')\n",
    "ylabel('|S21| (dB)')\n",
    "xlabel('MHz')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
